package com.smart.milvus.bean;

import com.alibaba.fastjson.JSONObject;
import com.google.gson.Gson;
import com.smart.milvus.base.FieldWrapperHandler;
import com.smart.milvus.base.JsonTypeHandler;
import com.smart.milvus.conditions.Wrapper;
import com.smart.milvus.conditions.query.BaseQueryWrapper;
import com.smart.milvus.conditions.query.QueryWrapper;
import com.smart.milvus.conditions.search.BaseSearchWrapper;
import com.smart.milvus.exception.MilvusException;
import com.smart.milvus.model.Page;
import com.smart.milvus.model.index.IIndexExtra;
import com.smart.milvus.model.metadata.CollectionDefinition;
import com.smart.milvus.model.metadata.CollectionHelper;
import com.smart.milvus.model.metadata.ColumnDefinition;
import com.smart.milvus.util.CollectionUtils;
import io.milvus.client.MilvusClient;
import io.milvus.grpc.*;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.*;
import io.milvus.param.dml.DeleteParam;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.QueryParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.param.index.DropIndexParam;
import io.milvus.response.QueryResultsWrapper;
import io.milvus.response.SearchResultsWrapper;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.BeanWrapper;
import org.springframework.beans.BeanWrapperImpl;
import org.springframework.util.ReflectionUtils;

import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;

@Slf4j
public class MilvusClientService {
    @Getter
    private final MilvusClient milvusClient;

    private final Gson gson = new Gson();

    public MilvusClientService(MilvusClient milvusClient) {
        this.milvusClient = milvusClient;
    }

    public <T extends MilvusModel<?>> boolean remove(Object pk, Class<T> clazz) throws MilvusException {
        CollectionDefinition collection = CollectionHelper.getCollectionInfo(clazz);
        String columnName = collection.getPrimaryColumn().getName();
        String expression = new QueryWrapper<T>().in(columnName, pk).getExprSegment();
        return remove(collection, expression);
    }

    private boolean remove(CollectionDefinition collection, String expression) {
        if (StringUtils.isEmpty(expression)) {
            throw new MilvusException("expression is null");
        }
        DeleteParam.Builder builder = DeleteParam.newBuilder().withCollectionName(collection.getName()).withExpr(expression);
        R<MutationResult> resultR = milvusClient.delete(builder.build());
        if (resultR.getStatus() != R.Status.Success.getCode() || resultR.getException() != null) {
            throw new MilvusException(resultR.getException().getMessage());
        }
        return true;
    }

    public <T extends MilvusModel<?>> boolean batchRemove(Wrapper<T> wrapper) throws MilvusException {
        CollectionDefinition collection = CollectionHelper.getCollectionInfo(wrapper.getEntityClass());
        String expression = wrapper.getTargetExpr();
        return remove(collection, expression);
    }

    public <T extends MilvusModel<?>> Boolean insert(T vectorModel) throws MilvusException {
        CollectionDefinition collectionDefinition = CollectionHelper.getCollectionInfo(vectorModel.getClass());
        InsertParam.Builder builder = InsertParam.newBuilder();
        List<InsertParam.Field> dataFields = new ArrayList<>();
        BeanWrapper beanWrapper = new BeanWrapperImpl(vectorModel);
        for (ColumnDefinition columnDefinition : collectionDefinition.getColumns()) {
            Object value = beanWrapper.getPropertyValue(columnDefinition.getField().getName());
            String columnName = columnDefinition.getName();
            if (value == null) {
                continue;
            }

            if (columnDefinition.getTypeHandler() != null) {
                value = columnDefinition.getTypeHandler().replaceParam(value, columnDefinition.getField(), columnName, columnDefinition.getDataType(), columnDefinition.getElementType());
            } else {
                if (columnDefinition.isJsonField()) {
                    JsonTypeHandler jsonTypeHandler = GlobalConfig.INSTANCE.getJsonTypeHandler();
                    value = jsonTypeHandler.toJson(value);
                }
                if (columnDefinition.isJsonArray()) {
                    JsonTypeHandler jsonTypeHandler = GlobalConfig.INSTANCE.getJsonTypeHandler();
                    List<JSONObject> list = new ArrayList<>();
                    if (columnDefinition.getField().getType().isArray()) {
                        for (Object o : Arrays.asList(value)) {
                            list.add(jsonTypeHandler.toJson(o));
                        }
                    } else if (Collection.class.isAssignableFrom(columnDefinition.getField().getType())) {
                        for (Object o : ((Collection) value)) {
                            list.add(jsonTypeHandler.toJson(o));
                        }
                    }
                    value = list;
                }
            }

            FieldWrapperHandler<Object> vectorTypeHandler = columnDefinition.getFieldWrapperHandler();
            List<Object> dataVector = vectorTypeHandler.wrap(value);
            dataFields.add(new InsertParam.Field(columnName, dataVector));
        }
        if (!StringUtils.isEmpty(collectionDefinition.getDatabase())) {
            builder.withDatabaseName(collectionDefinition.getDatabase());
        }
        builder.withCollectionName(collectionDefinition.getName());
        builder.withFields(dataFields);
        InsertParam insertParam = builder.build();
        R<MutationResult> resultR = milvusClient.insert(insertParam);
        if (resultR.getStatus() != R.Status.Success.getCode() || resultR.getException() != null) {
            throw new MilvusException(resultR.getException().getMessage());
        }
        if (resultR.getData() != null && !resultR.getData().getIDs().getIntId().getDataList().isEmpty()) {
            ColumnDefinition column = collectionDefinition.getPrimaryColumn();
            ReflectionUtils.makeAccessible(column.getField());
            Object id = resultR.getData().getIDs().getIntId().getDataList().get(0);
            ReflectionUtils.setField(column.getField(), vectorModel, id);
        }
        return true;
    }

    public <T extends MilvusModel<?>> void loadCollection(Class<T> clazz) throws MilvusException {
        CollectionDefinition collectionDefinition = CollectionHelper.getCollectionInfo(clazz);
        LoadCollectionParam.Builder builder = LoadCollectionParam.newBuilder();
        builder.withCollectionName(collectionDefinition.getName());
        if (!StringUtils.isEmpty(collectionDefinition.getDatabase())) {
            builder.withDatabaseName(collectionDefinition.getDatabase());
        }
        R<RpcStatus> resultR = milvusClient.loadCollection(builder.build());
        if (resultR.getStatus() != R.Status.Success.getCode() || resultR.getException() != null) {
            throw new MilvusException(resultR.getException().getMessage());
        }
    }

    public <T extends MilvusModel<?>> void releaseCollection(Class<T> clazz) throws MilvusException {
        CollectionDefinition collectionDefinition = CollectionHelper.getCollectionInfo(clazz);
        milvusClient.releaseCollection(ReleaseCollectionParam.newBuilder().withCollectionName(collectionDefinition.getName()).build());
    }

    public <T extends MilvusModel<?>> void dropCollection(Class<T> clazz) throws MilvusException {
        CollectionDefinition collectionDefinition = CollectionHelper.getCollectionInfo(clazz);
        DropCollectionParam.Builder builder = DropCollectionParam.newBuilder().withCollectionName(collectionDefinition.getName());
        if (!StringUtils.isEmpty(collectionDefinition.getDatabase())) {
            builder.withDatabaseName(collectionDefinition.getDatabase());
        }
        R<RpcStatus> resultR = milvusClient.dropCollection(builder.build());
        if (resultR.getStatus() != R.Status.Success.getCode() || resultR.getException() != null) {
            throw new MilvusException(resultR.getException().getMessage());
        }
    }

    public <T extends MilvusModel<?>> boolean hasCollection(Class<T> clazz) throws MilvusException {
        CollectionDefinition collectionDefinition = CollectionHelper.getCollectionInfo(clazz);
        HasCollectionParam.Builder builder = HasCollectionParam.newBuilder().withCollectionName(collectionDefinition.getName());
        if (!StringUtils.isEmpty(collectionDefinition.getDatabase())) {
            builder.withDatabaseName(collectionDefinition.getDatabase());
        }
        R<Boolean> resultR = milvusClient.hasCollection(builder.build());
        if (resultR.getStatus() != R.Status.Success.getCode() || resultR.getException() != null) {
            throw new MilvusException(resultR.getException().getMessage());
        }
        return resultR.getData();
    }

    public <T extends MilvusModel<?>> boolean dropIndex(Class<T> clazz, String indexName) throws MilvusException {
        CollectionDefinition collectionDefinition = CollectionHelper.getCollectionInfo(clazz);
        DropIndexParam.Builder builder = DropIndexParam.newBuilder().withCollectionName(collectionDefinition.getName()).withIndexName(indexName);
        R<RpcStatus> resultR = milvusClient.dropIndex(builder.build());
        if (resultR.getStatus() != R.Status.Success.getCode() || resultR.getException() != null) {
            throw new MilvusException(resultR.getException().getMessage());
        }
        return true;
    }

    public <T extends MilvusModel<?>> LoadState getLoadState(Class<T> clazz) throws MilvusException {
        CollectionDefinition collectionDefinition = CollectionHelper.getCollectionInfo(clazz);
        GetLoadStateParam.Builder builder = GetLoadStateParam.newBuilder();
        builder.withCollectionName(collectionDefinition.getName());
        if (!StringUtils.isEmpty(collectionDefinition.getDatabase())) {
            builder.withDatabaseName(collectionDefinition.getDatabase());
        }
        R<GetLoadStateResponse> resultR = milvusClient.getLoadState(builder.build());
        if (resultR.getStatus() != R.Status.Success.getCode() || resultR.getException() != null) {
            throw new MilvusException(resultR.getException().getMessage());
        }
        return resultR.getData().getState();
    }

    public <T extends MilvusModel<?>> Long getLoadProgress(Class<T> clazz) throws MilvusException {
        CollectionDefinition collectionDefinition = CollectionHelper.getCollectionInfo(clazz);
        GetLoadingProgressParam.Builder builder = GetLoadingProgressParam.newBuilder();
        builder.withCollectionName(collectionDefinition.getName());
        R<GetLoadingProgressResponse> resultR = milvusClient.getLoadingProgress(builder.build());
        if (resultR.getStatus() != R.Status.Success.getCode() || resultR.getException() != null) {
            throw new MilvusException(resultR.getException().getMessage());
        }
        return resultR.getData().getProgress();
    }

    public <T extends MilvusModel<?>> boolean createIndex(BaseSearchWrapper<T, ?, ? extends BaseSearchWrapper<T, ?, ?>> wrapper) throws MilvusException {
        return createIndex(wrapper.getEntityClass(), wrapper.getIndexName(), wrapper.getIndexColumn(), wrapper.getExtra());
    }

    public <T extends MilvusModel<?>> boolean createIndex(Class<T> clazz, String indexName, String columnName, IIndexExtra extra) throws MilvusException {
        CollectionDefinition collectionDefinition = CollectionHelper.getCollectionInfo(clazz);
        CreateIndexParam.Builder builder = CreateIndexParam.newBuilder();
        builder.withCollectionName(collectionDefinition.getName());
        ColumnDefinition columnDefinition = collectionDefinition.getColumnByColumnName(columnName);
        builder.withFieldName(columnName);
        builder.withIndexName(indexName);
        builder.withIndexType(columnDefinition.getIndexType());
        builder.withMetricType(columnDefinition.getMetricType());
        if (extra != null) {
            builder.withExtraParam(new Gson().toJson(extra));
        }
        builder.withSyncMode(Boolean.FALSE);
        R<RpcStatus> resultR = milvusClient.createIndex(builder.build());
        if (resultR.getStatus() != R.Status.Success.getCode() || resultR.getException() != null) {
            throw new MilvusException(resultR.getException().getMessage());
        }
        return true;
    }

    public <T extends MilvusModel<?>> boolean createCollection(Class<T> clazz) throws MilvusException {
        CollectionDefinition collectionDefinition = CollectionHelper.getCollectionInfo(clazz);
        CreateCollectionParam.Builder builder = CreateCollectionParam.newBuilder();
        builder.withCollectionName(collectionDefinition.getName());
        builder.withDescription(collectionDefinition.getDescription());
        for (ColumnDefinition column : collectionDefinition.getColumns()) {
            FieldType.Builder fieldBuilder = FieldType.newBuilder();
            fieldBuilder.withName(column.getName());
            fieldBuilder.withDescription(column.getDesc());
            fieldBuilder.withPartitionKey(column.isPartitionKey());
            fieldBuilder.withDataType(column.getDataType());
            Boolean primary = column.isPrimary();
            fieldBuilder.withPrimaryKey(primary);
            if (column.vectorColumn()) {
                fieldBuilder.withDimension(column.getVectorDimension());
            }
            if (column.getDataType() == DataType.VarChar) {
                fieldBuilder.withMaxLength(column.getMaxLength());
            }
            if (column.getDataType() == DataType.Array) {
                if (column.getElementType() == DataType.None) {
                    throw new MilvusException(column.getName() + " column must set elementType");
                }
                fieldBuilder.withElementType(column.getElementType());
                fieldBuilder.withMaxCapacity(column.getMaxCapacity());
                fieldBuilder.withMaxLength(column.getMaxLength());
            }
            if (Boolean.TRUE.equals(primary)) {
                fieldBuilder.withAutoID(true);
            }
            if (column.canBePartitionKey()) {
                fieldBuilder.withPartitionKey(column.isPartitionKey());
            }
            builder.addFieldType(fieldBuilder.build());
        }
        R<RpcStatus> resultR = milvusClient.createCollection(builder.build());
        if (resultR.getStatus() != R.Status.Success.getCode() || resultR.getException() != null) {
            throw new MilvusException(resultR.getException().getMessage());
        }
        return true;
    }

    public <T extends MilvusModel<?>> List<T> search(BaseSearchWrapper<T, ?, ? extends BaseSearchWrapper<T, ?, ?>> wrapper) throws MilvusException {
        return this.search(wrapper, wrapper.getEntityClass());
    }

    public <T extends MilvusModel<?>> List<T> search(BaseSearchWrapper<T, ?, ? extends BaseSearchWrapper<T, ?, ?>> wrapper, Class<T> clazz) throws MilvusException {
        CollectionDefinition collectionDefinition = CollectionHelper.getCollectionInfo(clazz);
        List<String> outFields = new ArrayList<>();
        for (ColumnDefinition columnDefinition : collectionDefinition.getColumns()) {
            if (columnDefinition.vectorColumn()) {
                continue;
            }
            outFields.add(columnDefinition.getName());
        }
        SearchParam.Builder builder = SearchParam.newBuilder();
        String vectorColumnName = wrapper.getVectorColumn();
        ColumnDefinition columnDefinition = collectionDefinition.getColumnByColumnName(vectorColumnName);
        FieldWrapperHandler<Object> vectorTypeHandler = columnDefinition.getFieldWrapperHandler();
        List<?> vectors = vectorTypeHandler.wrap(wrapper.getVectorValue());
        builder.withVectors(vectors);
        builder.withVectorFieldName(columnDefinition.getName());
        builder.withCollectionName(collectionDefinition.getName());
        builder.withConsistencyLevel(wrapper.getConsistencyLevel());
        builder.withMetricType(columnDefinition.getMetricType());
        builder.withOutFields(outFields);
        builder.withTopK(wrapper.getTopK());
        String expression = wrapper.getExprSegment();
        if (!StringUtils.isEmpty(expression)) {
            builder.withExpr(expression);
        }
        if (!CollectionUtils.isEmpty(wrapper.getPartitionNames())) {
            builder.withPartitionNames(wrapper.getPartitionNames());
        }
        if (wrapper.getExtra() != null) {
            builder.withParams(gson.toJson(wrapper.getExtra()));
        }
        R<SearchResults> resultR = milvusClient.search(builder.build());
        if (resultR.getStatus() != R.Status.Success.getCode() || resultR.getException() != null) {
            throw new MilvusException(resultR.getException().getMessage());
        }
        SearchResultsWrapper resultsWrapper = new SearchResultsWrapper(resultR.getData().getResults());

        return fillProp(clazz, resultsWrapper.getRowRecords());
    }

    private <T extends MilvusModel<?>> List<T> fillProp(Class<T> clazz, List<QueryResultsWrapper.RowRecord> rowRecords) {
        CollectionDefinition collectionDefinition = CollectionHelper.getCollectionInfo(clazz);
        JsonTypeHandler jsonTypeHandler = GlobalConfig.INSTANCE.getJsonTypeHandler();
        List<T> resultRows = new ArrayList<>();
        for (int i = 0; i < rowRecords.size(); i++) {
            QueryResultsWrapper.RowRecord rowRecord = rowRecords.get(i);
            BeanWrapper beanWrapper = new BeanWrapperImpl(clazz);
            rowRecord.getFieldValues().keySet().forEach(columnName -> {
                ColumnDefinition column = collectionDefinition.getColumnByColumnName(columnName);
                if (column == null) {
                    return;
                }
                Object value = rowRecord.get(columnName);

                if (column.getTypeHandler() != null) {
                    value = column.getTypeHandler().getResult(value, column.getField(), columnName, column.getDataType(), column.getElementType());
                } else {
                    if (column.isJsonField() && value instanceof JSONObject) {
                        value = jsonTypeHandler.parse((JSONObject) value, column.getField().getType());
                    }
                    if (column.isJsonArray() && Collection.class.isAssignableFrom(column.getField().getType()) && value instanceof Collection) {
                        List<Object> list = new ArrayList<>();
                        for (Object o : ((Collection) value)) {
                            if (o instanceof JSONObject) {
                                Type genericType = column.getField().getGenericType();
                                Class targetType = Object.class;
                                if (genericType instanceof ParameterizedType) {
                                    ParameterizedType genericType1 = (ParameterizedType) genericType;
                                    targetType = (Class) genericType1.getActualTypeArguments()[0];
                                }
                                list.add(jsonTypeHandler.parse((JSONObject) o, targetType));
                            }
                        }
                        value = list;
                    }
                }

                beanWrapper.setPropertyValue(column.getField().getName(), value);
            });
            resultRows.add((T) beanWrapper.getWrappedInstance());
        }
        return resultRows;
    }

    public <T extends MilvusModel<?>> List<T> query(BaseQueryWrapper<T, ? extends BaseQueryWrapper<T, ?>> wrapper) throws MilvusException {
        Class<T> clazz = wrapper.getEntityClass();
        CollectionDefinition collectionDefinition = CollectionHelper.getCollectionInfo(clazz);
        List<String> outFields = new ArrayList<>();
        for (ColumnDefinition columnDefinition : collectionDefinition.getColumns()) {
            if (columnDefinition.vectorColumn()) {
                continue;
            }
            outFields.add(columnDefinition.getName());
        }
        QueryParam.Builder builder = QueryParam.newBuilder();
        if (!CollectionUtils.isEmpty(wrapper.getPartitionNames())) {
            builder.withPartitionNames(wrapper.getPartitionNames());
        }
        if (!CollectionUtils.isEmpty(wrapper.getPartitionNames())) {
            builder.withPartitionNames(wrapper.getPartitionNames());
        }
        builder.withCollectionName(collectionDefinition.getName());
        builder.withConsistencyLevel(wrapper.getConsistencyLevel());
        builder.withOutFields(outFields);
        if (wrapper.getLimit() != null) {
            builder.withLimit(wrapper.getLimit());
        }
        if (wrapper.getOffset() != null) {
            builder.withOffset(wrapper.getOffset());
        }
        String expression = wrapper.getTargetExpr();
        if (!StringUtils.isEmpty(expression)) {
            builder.withExpr(expression);
        }
        R<QueryResults> resultR = milvusClient.query(builder.build());
        if (resultR.getStatus() != R.Status.Success.getCode() || resultR.getException() != null) {
            throw new MilvusException(resultR.getException().getMessage());
        }
        QueryResultsWrapper resultsWrapper = new QueryResultsWrapper(resultR.getData());
        return fillProp(clazz, resultsWrapper.getRowRecords());
    }

    public <T extends MilvusModel<?>> GetCollectionStatisticsResponse getCollectionStatistics(Class<T> clazz) throws MilvusException {
        CollectionDefinition collectionDefinition = CollectionHelper.getCollectionInfo(clazz);
        GetCollectionStatisticsParam.Builder builder = GetCollectionStatisticsParam.newBuilder();
        builder.withCollectionName(collectionDefinition.getName());
        if (!StringUtils.isEmpty(collectionDefinition.getDatabase())) {
            builder.withDatabaseName(collectionDefinition.getDatabase());
        }
        R<GetCollectionStatisticsResponse> resultR = milvusClient.getCollectionStatistics(builder.build());
        if (resultR.getStatus() != R.Status.Success.getCode() || resultR.getException() != null) {
            throw new MilvusException(resultR.getException().getMessage());
        }
        return resultR.getData();
    }

    public <T extends MilvusModel<?>> Page<T> queryPage(BaseQueryWrapper<T, ? extends BaseQueryWrapper<T, ?>> wrapper, Long page, Long pageSize) throws MilvusException {
        Page<T> dataPage = new Page<>();
        dataPage.setPage(page);
        dataPage.setPageSize(pageSize);
        wrapper.setLimit(pageSize);
        wrapper.setOffset(page * pageSize);
        List<T> instanceList = this.query(wrapper);
        dataPage.setInstances(instanceList);
        return dataPage;
    }

    public <T extends MilvusModel<?>> Long getRowCount(Class<T> clazz) throws MilvusException {
        GetCollectionStatisticsResponse statistics = getCollectionStatistics(clazz);
        for (KeyValuePair pair : statistics.getStatsList()) {
            if ("row_count".equals(pair.getKey())) {
                return Long.parseLong(pair.getValue());
            }
        }
        return 0L;
    }
}
